/*
 * OpenDPI_demo.c
 * Copyright (C) 2009-2011 by ipoque GmbH
 * 
 * This file is part of OpenDPI, an open source deep packet inspection
 * library based on the PACE technology by ipoque GmbH
 * 
 * OpenDPI is free software: you can redistribute it and/or modify
 * it under the terms of the GNU Lesser General Public License as published by
 * the Free Software Foundation, either version 3 of the License, or
 * (at your option) any later version.
 * 
 * OpenDPI is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 * GNU Lesser General Public License for more details.
 * 
 * You should have received a copy of the GNU Lesser General Public License
 * along with OpenDPI.  If not, see <http://www.gnu.org/licenses/>.
 * 
 */



#include <stdio.h>
#include <stdlib.h>
#include <unistd.h>
#include <string.h>
#include <stdarg.h>
#include <netinet/in.h>

#ifdef __linux__
#include <linux/ip.h>
#include <netinet/ip6.h>
#include <linux/tcp.h>
#include <linux/udp.h>
#include <mysql/mysql.h>
#include <linux/if_ether.h>
#else
# include "linux_compat.h"
#endif

#include <pcap.h>
#include <assert.h>

#include "ipq_api.h"
#include "OpenDPI_demo.h"
#include "sklcc_api.h"

// cli options
static char *_pcap_file = NULL;

//flag for flow
static u16 _flow_flag = 0;
// pcap
char _pcap_error_buffer[PCAP_ERRBUF_SIZE];
static pcap_t *_pcap_handle = NULL;
static int _pcap_datalink_type = 0;

// detection
static struct ipoque_detection_module_struct *ipoque_struct = NULL;
static u32 detection_tick_resolution = 1000;
char *prot_long_str[] = { IPOQUE_PROTOCOL_LONG_STRING };

#ifdef IPOQUE_ENABLE_DEBUG_MESSAGES
char *prot_short_str[] = { IPOQUE_PROTOCOL_SHORT_STRING };

IPOQUE_PROTOCOL_BITMASK debug_messages_bitmask;
#endif

// results
static u64 raw_packet_count = 0;
u64 ip_packet_count = 0;
u64 total_bytes = 0;
u64 protocol_counter[IPOQUE_MAX_SUPPORTED_PROTOCOLS + 1];
u64 protocol_counter_bytes[IPOQUE_MAX_SUPPORTED_PROTOCOLS + 1];

u32 size_id_struct = 0;
struct osdpi_id *osdpi_ids;
u32 osdpi_id_count = 0;

u32 size_flow_struct = 0;
struct osdpi_flow *osdpi_flows;
u32 osdpi_flow_count = 0;

#ifdef IPOQUE_ENABLE_DEBUG_MESSAGES
static int string_to_detection_bitmask(char *str, IPOQUE_PROTOCOL_BITMASK * dbm)
{
	u32 a;
	u32 oldptr = 0;
	u32 ptr = 0;
	IPOQUE_BITMASK_RESET(*dbm);

	printf("Protocol parameter given: %s\n", str);

	if (strcmp(str, "all") == 0) {
		printf("Protocol parameter all parsed\n");
		IPOQUE_BITMASK_SET_ALL(*dbm);
		printf("Bitmask is: " IPOQUE_BITMASK_DEBUG_OUTPUT_BITMASK_STRING " \n",
			   IPOQUE_BITMASK_DEBUG_OUTPUT_BITMASK_VALUE(*dbm));
		return 0;
	}
	// parse bitmask
	while (1) {
		if (str[ptr] == 0 || str[ptr] == ' ') {
			printf("Protocol parameter: parsed: %.*s,\n", ptr - oldptr, &str[oldptr]);
			for (a = 1; a <= IPOQUE_MAX_SUPPORTED_PROTOCOLS; a++) {

				if (strlen(prot_short_str[a]) == (ptr - oldptr) &&
					(memcmp(&str[oldptr], prot_short_str[a], ptr - oldptr) == 0)) {
					IPOQUE_ADD_PROTOCOL_TO_BITMASK(*dbm, a);
					printf("Protocol parameter detected as protocol %s\n", prot_long_str[a]);
				}
			}
			oldptr = ptr + 1;
			if (str[ptr] == 0)
				break;
		}
		ptr++;
	}
	return 0;
}
#endif

static void parseOptions(int argc, char **argv)
{
	int opt;

#ifdef IPOQUE_ENABLE_DEBUG_MESSAGES
	IPOQUE_BITMASK_SET_ALL(debug_messages_bitmask);
#endif
	if(argc == 1)
	{
		printf("Choose Live Mode\n");
		return ;
	}

	while ((opt = getopt(argc, argv, "f:e:i")) != EOF) {
		switch (opt) {
		case 'i':
			return ;
		case 'f':
			printf("Choose File Mode\n");
			_pcap_file = optarg;
			break;
		case 'e':
#ifdef IPOQUE_ENABLE_DEBUG_MESSAGES
			// set debug logging bitmask to all protocols
			if (string_to_detection_bitmask(optarg, &debug_messages_bitmask) != 0) {
				printf("ERROR option -e needs a valid list of protocols");
				exit(-1);
			}

			printf("debug messages Bitmask is: " IPOQUE_BITMASK_DEBUG_OUTPUT_BITMASK_STRING "\n",
				   IPOQUE_BITMASK_DEBUG_OUTPUT_BITMASK_VALUE(debug_messages_bitmask));

#else
			printf("ERROR: option -e : DEBUG MESSAGES DEACTIVATED\n");
			exit(-1);
#endif
			break;
		}
	}

	// check parameters
	if (_pcap_file == NULL || strcmp(_pcap_file, "") == 0) {
		printf("ERROR: no pcap file path provided; use option -f with the path to a valid pcap file\n");
		exit(-1);
	}
}

static void debug_printf(u32 protocol, void *id_struct, ipq_log_level_t log_level, const char *format, ...)
{
#ifdef IPOQUE_ENABLE_DEBUG_MESSAGES
	if (IPOQUE_COMPARE_PROTOCOL_TO_BITMASK(debug_messages_bitmask, protocol) != 0) {
		const char *protocol_string;
		const char *file;
		const char *func;
		u32 line;
		va_list ap;
		va_start(ap, format);

		protocol_string = prot_short_str[protocol];

		ipoque_debug_get_last_log_function_line(ipoque_struct, &file, &func, &line);

		printf("\nDEBUG: %s:%s:%u Prot: %s, level: %u packet: %llu :", file, func, line, protocol_string,
			   log_level, raw_packet_count);
		vprintf(format, ap);
		va_end(ap);
	}
#endif
}

static void *malloc_wrapper(unsigned long size)
{
	return malloc(size);
}

static void free_wrapper(void *freeable)
{
	free(freeable);
}

extern void *get_id(const struct osdpi_ip_addr *ip)
{
	u32 i;

	for (i = 0; i < osdpi_id_count; i++) {
		if (memcmp(&osdpi_ids[i].ip, ip, sizeof(struct osdpi_ip_addr)) == 0) {
			return osdpi_ids[i].ipoque_id;
		}
	}
	if (osdpi_id_count == MAX_OSDPI_IDS) {
		printf("ERROR: maximum unique id count (%u) has been exceeded\n", MAX_OSDPI_IDS);
		exit(-1);
	} else {
		struct ipoque_id_struct *ipoque_id;
		osdpi_ids[osdpi_id_count].ip = *ip;
		ipoque_id = osdpi_ids[osdpi_id_count].ipoque_id;

		osdpi_id_count += 1;
		return ipoque_id;
	}
}

extern struct osdpi_flow *get_osdpi_flow(const struct iphdr *iph, u16 ipsize)
{
	u32 i;
	const u8 *l4 = NULL;
	u16 l4_len = 0;
	u8 l4_protocol = 0;
	struct ipoque_unique_flow_ipv4_and_6_struct key;
	u8 dir = 0;

	if (ipoque_detection_get_l4((u8*)iph, ipsize, &l4, &l4_len, &l4_protocol, 0) == 0) {
		if (ipoque_detection_build_key((u8*)iph, ipsize, l4, l4_len, l4_protocol, &key, &dir, 0) == 0) {

			for (i = 0; i < osdpi_flow_count; i++) {
				if (memcmp(&osdpi_flows[i].key, &key, sizeof(struct ipoque_unique_flow_ipv4_and_6_struct)) == 0) {
					_flow_flag = 0;//old flow
					return &osdpi_flows[i];
				}
			}
			if (osdpi_flow_count == MAX_OSDPI_FLOWS) {
				printf("ERROR: maximum flow count (%u) has been exceeded\n", MAX_OSDPI_FLOWS);
				exit(-1);
			} else {
				struct osdpi_flow *flow;
				osdpi_flows[osdpi_flow_count].key = key;
				flow = &osdpi_flows[osdpi_flow_count];

				osdpi_flow_count += 1;
				_flow_flag = 1;//new flow
				return flow;
			}

		}
	}

	return NULL;
}

extern void setupDetection(void)
{
	printf("setup detection\n");
	PRINT_AUTHOR_INFO("ad_init","YangHao");
	ad_init();

	u32 i;
	IPOQUE_PROTOCOL_BITMASK all;

	// init global detection structure
	ipoque_struct = ipoque_init_detection_module(detection_tick_resolution, malloc_wrapper, debug_printf);
	if (ipoque_struct == NULL) {
		printf("ERROR: global structure initialization failed\n");
		exit(-1);
	}
	// enable all protocols
	IPOQUE_BITMASK_SET_ALL(all);
	ipoque_set_protocol_detection_bitmask2(ipoque_struct, &all);

	// allocate memory for id and flow tracking
	size_id_struct = ipoque_detection_get_sizeof_ipoque_id_struct();
	size_flow_struct = ipoque_detection_get_sizeof_ipoque_flow_struct();

	PRINT_AUTHOR_INFO("memop_osdpi_ids_create","ChenYaDong");
	osdpi_ids = memop_osdpi_ids_create(MAX_OSDPI_IDS);
	assert(osdpi_ids != NULL);	

	PRINT_AUTHOR_INFO("memop_osdpi_flows_create","ChenYaDong");
	osdpi_flows = memop_osdpi_flows_create(MAX_OSDPI_FLOWS); 
	assert(osdpi_flows != NULL);

	// clear memory for results
	memset(protocol_counter, 0, (IPOQUE_MAX_SUPPORTED_PROTOCOLS + 1) * sizeof(u64));
	memset(protocol_counter_bytes, 0, (IPOQUE_MAX_SUPPORTED_PROTOCOLS + 1) * sizeof(u64));
}

static void terminateDetection(void)
{
	u32 i;

	ipoque_exit_detection_module(ipoque_struct, free_wrapper);

	for (i = 0; i < MAX_OSDPI_IDS; i++) {
		free(osdpi_ids[i].ipoque_id);
	}
	free(osdpi_ids);
	for (i = 0; i < MAX_OSDPI_FLOWS; i++) {
		free(osdpi_flows[i].ipoque_flow);
	}
	free(osdpi_flows);
}

extern unsigned int packet_processing(int *protocol,const struct iphdr *iph,uint16_t ipsize,uint16_t rawsize)
{
	//BEGIN
	id_and_flow_struct id_and_flow;
	id_and_flow.osdpi_id_count = osdpi_id_count;
	id_and_flow.osdpi_flow_count = osdpi_flow_count;
	id_and_flow.osdpi_ids = osdpi_ids;
	id_and_flow.osdpi_flows = osdpi_flows;

	PRINT_AUTHOR_INFO("memop_action_start","ChengYaDong");
	memop_action_start((void *)&(id_and_flow));

	osdpi_ids = id_and_flow.osdpi_ids;
	osdpi_flows = id_and_flow.osdpi_flows;
	osdpi_id_count = id_and_flow.osdpi_id_count;
	osdpi_flow_count = id_and_flow.osdpi_flow_count;
	//END

	struct ipoque_id_struct *src = NULL;
	struct ipoque_id_struct *dst = NULL;
	struct osdpi_flow *flow = NULL;
	struct ipoque_flow_struct *ipq_flow = NULL;
	int proto = 0;

	struct osdpi_ip_addr src_ip, dst_ip;
	memset(&src_ip, 0, sizeof(struct osdpi_ip_addr));
	memset(&dst_ip, 0, sizeof(struct osdpi_ip_addr));

#ifdef IPOQUE_DETECTION_SUPPORT_IPV6
	if (iph->version == 6 && ipsize >= sizeof(struct ip6_hdr)) {
		struct ip6_hdr *ip6h = (struct ip6_hdr *)iph;
		src_ip.is_ip_v6 = 1;
		src_ip.ipv6 = ip6h->ip6_src;
		dst_ip.is_ip_v6 = 1;
		dst_ip.ipv6 = ip6h->ip6_dst;
	} else
#endif
	if (iph->version == 4 && ipsize >= sizeof(struct iphdr)) {
		src_ip.ipv4 = iph->saddr;
		dst_ip.ipv4 = iph->daddr;
	} else {
		return 1;
	}

	PRINT_AUTHOR_INFO("memop_get_id","ChenYaDong");
	src = memop_get_id(&src_ip);
	PRINT_AUTHOR_INFO("memop_get_id","ChenYaDong");
	dst = memop_get_id(&dst_ip);

	PRINT_AUTHOR_INFO("memop_get_osdpi_flow","ChenYaDong");
	flow = memop_get_osdpi_flow(iph, ipsize);
	if (flow != NULL) {
		ipq_flow = flow->ipoque_flow;
	} else {
		return 1;
	}

	ip_packet_count++;
	total_bytes += rawsize;

#ifndef IPOQUE_ENABLE_DEBUG_MESSAGES
	if (ip_packet_count % 499 == 0) {
		printf("\rip packets scanned: \x1b[33m%-10llu\x1b[0m ip bytes scanned: \x1b[34m%-10llu\x1b[0m",
			   ip_packet_count, total_bytes);
	}
#endif

	// only handle unfragmented packets
	if (iph->version == 4 && (iph->frag_off & htons(0x1FFF)) != 0) {

		static u8 frag_warning_used = 0;
		if (frag_warning_used == 0) {
			printf("\n\nWARNING: fragmented ip packets are not supported and will be skipped \n\n");
			sleep(2);
			frag_warning_used = 1;
		}
		return 0;

	} else {

		// here the actual detection is performed
		proto = ipoque_detection_process_packet(ipoque_struct, ipq_flow, (uint8_t *) iph, ipsize, time, src, dst);

	}
	
	protocol_counter[proto]++;
	protocol_counter_bytes[proto]+=rawsize;
	
	if (flow != NULL) {
		u16 flow_flag = _flow_flag;
		u32 ad_ret = 1;
		PRINT_AUTHOR_INFO("ad_process","YangHao");
		ad_ret = ad_process(flow_flag,flow,proto,iph,ipsize,rawsize);//for packets store
		assert(ad_ret == 0); 

		
		*protocol =proto;
		flow->detected_protocol = proto;

	}

	return 0;
}

extern void flow_statistics(struct flow_link *detected_flow_link,u32 de_protocol)
{
	struct packet_link *detected_packet = detected_flow_link->head_packet;
	while(detected_packet != NULL){
		protocol_counter[de_protocol]++;
		protocol_counter_bytes[de_protocol] += detected_packet->rawsize;
		detected_packet =detected_packet->next;
	}
	
}

extern void un_flow_statistics(u64 un_count,u64 un_count_byte)
{
	protocol_counter[0] += un_count;
	protocol_counter_bytes[0] += un_count_byte;
}

extern void printResults(void)
{
	u32 i;

	printf("\x1b[2K\n");
	printf("pcap file contains\n");
	printf("\tip packets:   \x1b[33m%-13llu\x1b[0m of %llu packets total\n", ip_packet_count, raw_packet_count);
	printf("\tip bytes:     \x1b[34m%-13llu\x1b[0m\n", total_bytes);
	printf("\tunique ids:   \x1b[35m%-13u\x1b[0m\n", osdpi_id_count);
	printf("\tunique flows: \x1b[36m%-13u\x1b[0m\n", osdpi_flow_count);

	printf("\n\ndetected protocols:\n");
	for (i = 0; i <= IPOQUE_MAX_SUPPORTED_PROTOCOLS; i++) {
		u32 protocol_flows = 0;
		u32 j;

		// count flows for that protocol
		for (j = 0; j < osdpi_flow_count; j++) {
			if (osdpi_flows[j].detected_protocol == i) {
				protocol_flows++;
			}
		}

		if (protocol_counter[i] > 0) {
			printf("\t\x1b[31m%-20s\x1b[0m packets: \x1b[33m%-13llu\x1b[0m bytes: \x1b[34m%-13llu\x1b[0m "
				   "flows: \x1b[36m%-13u\x1b[0m\n",
				   prot_long_str[i], protocol_counter[i], protocol_counter_bytes[i], protocol_flows);
		}
	}
	printf("\n\n");
}

static u32 table_protocol_id = 0;
extern void sync_sql(MYSQL *conn)
{
	MYSQL_RES *res_ptr;
	MYSQL_ROW sqlrow;
	int res;

	u32 i;
	for(i = 0; i < IPOQUE_MAX_SUPPORTED_PROTOCOLS; i++)
	{
		u32 protocol_flows = 0;
		u32 j;

		for(j = 0;j < osdpi_flow_count; j++)
		{
			if(osdpi_flows[j].detected_protocol == i)
			{
				protocol_flows++;
			}
		}

		if(protocol_counter[i] > 0)
		{
			char SQLcmd[1024];
			sprintf(SQLcmd,"insert into protocolStatistics2 values(%d,'%s',%lld,%lld,%d)",
					table_protocol_id,prot_long_str[i],protocol_counter[i],
					protocol_counter_bytes[i],protocol_flows);

			res = mysql_query(conn,SQLcmd);
			if(res)
			{
				printf("insert error %d : %s\n",mysql_errno(conn),mysql_error(conn));
			}
		}
	}

	table_protocol_id++;
}

static void openPcapFile()
{
	_pcap_handle = pcap_open_offline(_pcap_file,_pcap_error_buffer);

	if (_pcap_handle == NULL) {
		printf("ERROR: could not open pcap file: %s\n", _pcap_error_buffer);
		exit(-1);
	}
	_pcap_datalink_type = pcap_datalink(_pcap_handle);
}

static void openPcapDevice()
{
	char *dev_name = pcap_lookupdev(_pcap_error_buffer);
	if (dev_name == NULL) {
		printf("ERROR: could not find pcap device.Try it in root: %s\n", _pcap_error_buffer);
		exit(-1);
	}

	if(dev_name != NULL) {
		printf("Interface Name: %s\n",dev_name);
		_pcap_handle = pcap_open_live(dev_name,65536,1,1, _pcap_error_buffer);
	}

	if (_pcap_handle == NULL) {
		printf("ERROR: could not open pcap device.Try it in root: %s\n", _pcap_error_buffer);
		exit(-1);
	}
	_pcap_datalink_type = pcap_datalink(_pcap_handle);
}

static void closePcapFile(void)
{
	if (_pcap_handle != NULL) {
		pcap_close(_pcap_handle);
	}
}

static void print_time()
{
	struct tm *ptm; 
	long ts; 
	int y,m,d,h,n,s; 

	ts = time(NULL); 
	ptm = localtime(&ts); 

	y = ptm->tm_year+1900; //年 
	m = ptm->tm_mon+1; //月 
	d = ptm->tm_mday; //日 
	h = ptm->tm_hour; //时 
	n = ptm->tm_min; //分 
	s = ptm->tm_sec; //秒

	printf("%04d-%02d-%02d %02d:%02d:%02d\n",y,m,d,h,n,s);
}

static u32 i_am_alive = 0;
#define PRINT_FREQUENCE 1

static int am_i_alive()
{
	if(i_am_alive % PRINT_FREQUENCE == 0 && i_am_alive != 0)
	{
		printf("Get packet at : ");
		print_time();
		i_am_alive == 0;
		i_am_alive++;
		return 1;
	}
	else
	{
		i_am_alive++;
		return 0;
	}
}

/*
// executed for each packet in the pcap file
static void pcap_packet_callback(u_char * args, const struct pcap_pkthdr *header, const u_char * packet)
{
	am_i_alive();

	const struct ethhdr *ethernet = (struct ethhdr *) packet;
	struct iphdr *iph = (struct iphdr *) &packet[sizeof(struct ethhdr)];
	u64 time;
	static u64 lasttime = 0;
	u16 type;

	raw_packet_count++;

	time =
		((uint64_t) header->ts.tv_sec) * detection_tick_resolution +
		header->ts.tv_usec / (1000000 / detection_tick_resolution);
	if (lasttime > time) {
		// printf("\nWARNING: timestamp bug in the pcap file (ts delta: %llu, repairing)\n", lasttime - time);
		time = lasttime;
	}
	lasttime = time;


	type = ethernet->h_proto;

	// just work on Ethernet packets that contain IP
	if (_pcap_datalink_type == DLT_EN10MB &&
		(type == htons(ETH_P_IP)
#ifdef IPOQUE_DETECTION_SUPPORT_IPV6
		|| type == htons(ETH_P_IPV6)
#endif
		)
		&& header->caplen >= sizeof(struct ethhdr)) {

		if (header->caplen < header->len) {
			static u8 cap_warning_used = 0;
			if (cap_warning_used == 0) {
				printf
					("\n\nWARNING: packet capture size is smaller than packet size, DETECTION MIGHT NOT WORK CORRECTLY OR EVEN CRASH\n\n");
				sleep(2);
				cap_warning_used = 1;
			}
		}

		if (header->len >= (sizeof(struct ethhdr) + sizeof(struct iphdr))) {
			// process the packet
			packet_processing(time, iph, header->len - sizeof(struct ethhdr), header->len);
		}
	}
}

static void runPcapLoop(void)
{
	if (_pcap_handle != NULL) {
		pcap_loop(_pcap_handle, -1, &pcap_packet_callback, NULL);
	}
}

int main(int argc, char **argv)
{
	parseOptions(argc, argv);

	setupDetection();

	if(argc != 1)
		openPcapFile();
	else
		openPcapDevice();

	runPcapLoop();

	if(argc == 1)
		closePcapFile();

	printResults();

	terminateDetection();

	return 0;
}
*/
